import numpy as np
import matplotlib.pyplot as plt
from scipy.io import loadmat


def K_Means(K, dataset):
    # initialize K cluster centers
    mu = (np.random.random((K, 2)) + np.array([0, -0.5])) * np.array([10, 10])

    # iterate
    while True:
        # step 1: assign each sample to the nearest cluster center
        assign = np.array([np.argmin(np.linalg.norm(mu - x, axis=1)) for x in dataset])

        # step 2: update cluster center
        mu_new = np.array([np.sum(dataset[assign == i], axis=0) / np.sum(assign == i) for i in range(K)])

        # step 3: if coverage, end the iteration
        if np.linalg.norm(mu_new - mu) < 0.000001:
            return mu_new, assign
        mu = mu_new


def plot_cluster(dat, cent, dif, acc):
    shape = [',', 'o', 'v', '^', '<']
    for i in range(5):
        plt.scatter(dat[i * 200:i * 200 + 200][:, 0], dat[i * 200:i * 200 + 200][:, 1], marker=shape[i], s=20)
        plt.scatter(cent[i][0], cent[i][1], c='k', s=40)
    plt.title('Diff = {}, Accu = {}'.format(dif, acc))
    plt.show()


if __name__ == '__main__':
    data = np.array(loadmat('data/X.mat')['X'])
    cl_center = np.array(loadmat('data/mu.mat')['mu'])
    km_center, km_classification = K_Means(5, data)

    # calculate cluster id
    cl_id = [np.argmax(np.bincount(km_classification[i * 200:i * 200 + 200])) for i in range(5)]

    # difference from real cluster center to result cluster center
    diff = np.linalg.norm(cl_center - km_center[cl_id])

    # clustering accuracy
    accu = np.sum([np.sum(km_classification[i * 200:i * 200 + 200] == cl_id[i]) for i in range(5)]) / data.shape[0]

    print('...ok')
    plot_cluster(data, km_center[cl_id], diff, accu)
